import os
import pandas as pd
from torch.utils.data.dataset import Dataset
from torchvision.datasets.folder import default_loader
from torchvision.transforms import transforms
from torch.utils.data.dataloader import DataLoader

transform_train = transforms.Compose([
        transforms.Scale((550, 550)),
        transforms.RandomCrop(448, padding=8),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ])

transform_test = transforms.Compose([
        transforms.Scale((550, 550)),
        transforms.CenterCrop(448),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ])


class CUBDataSet(Dataset):

    def __init__(self, root, train=True):
        img_folder = os.path.join(root, "images")
        img_paths = pd.read_csv(os.path.join(root, "images.txt"), sep=" ", header=None, names=['idx', 'path'])
        img_labels = pd.read_csv(os.path.join(root, "image_class_labels.txt"), sep=" ", header=None,
                                 names=['idx', 'label'])
        train_test_split = pd.read_csv(os.path.join(root, "train_test_split.txt"), sep=" ", header=None,
                                       names=['idx', 'train_flag'])
        data = pd.concat([img_paths, img_labels, train_test_split], axis=1)
        data = data[data['train_flag'] == train]
        data['label'] = data['label'] - 1

        imgs = data.reset_index(drop=True)

        if len(imgs) == 0:
            raise (RuntimeError("no csv file"))
        self.root = img_folder
        self.imgs = imgs
        self.train = train

    def __getitem__(self, index):
        item = self.imgs.iloc[index]
        file_path = item['path']
        target = item['label']
        img = default_loader(os.path.join(self.root, file_path))
        if self.train:
            img = transform_train(img)
        else:
            img = transform_test(img)
        return img, target

    def __len__(self):
        return len(self.imgs)


if __name__ == '__main__':
    train_set = CUBDataSet("E:\PythonProj\BCNN\CUB_200_2011", train=True)
    train_loader = DataLoader(train_set, batch_size=5, shuffle=True)
    print(len(train_set))
    dataiter = iter(train_loader)
    images, labels = next(dataiter)
    print(images.shape)
    print(labels)
